import streamlit as st
from utils import dialogue_simulator
from langchain_community.document_loaders import PyPDFLoader, UnstructuredFileLoader


def main() -> None:
    st.title("GHAI Streamlit App")
    if "messages" not in st.session_state:
        st.session_state.messages = []

    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # upload pdf file
    if my_paper := st.file_uploader("Выберите документ", type=['txt', "pdf", 'md'], accept_multiple_files=False):
        if my_paper is not None:
            with st.chat_message("assistant"):
                temp_file = "./temp.pdf"
                with open(temp_file, "wb") as file:
                    file.write(my_paper.getvalue())
                    file_name = my_paper.name

                loader = UnstructuredFileLoader(temp_file, strategy="fast")
                pages = loader.load()
                data = ''
                for page in pages:
                    data += page.page_content

                topic, simulator = dialogue_simulator(data)
                n = 0
                max_iters = 5
                response = st.write('Тема для разговора: \n' + topic)
                st.session_state.messages.append({"role": 'Moderator', "content": response})
                while n < max_iters:
                    name, message = simulator.step()
                    print(f"({name}): {message}")
                    print("\n")
                    n += 1
                    response = st.write(f'{name}: {message}')
                    st.session_state.messages.append({"role": name, "content": response})
                simulator.reset()

    if user_input := st.text_input("Введите текст"):
        with st.chat_message("assistant"):
            topic, simulator = dialogue_simulator(user_input)
            n = 0
            max_iters = 5
            response = st.write('Тема для разговора: \n' + topic)
            st.session_state.messages.append({"role": 'Moderator', "content": response})
            while n < max_iters:
                name, message = simulator.step()
                print(f"({name}): {message}")
                print("\n")
                n += 1


if __name__ == "__main__":
    main()
